{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Gesture Classification using TAO GestureNet\n",
    "\n",
    "Transfer learning is the process of transferring learned features from one application to another. It is a commonly used training technique where you use a model trained on one task and re-train to use it on a different task. \n",
    "\n",
    "Train Adapt Optimize (TAO) Toolkit is a simple and easy-to-use Python based AI toolkit for taking purpose-built AI models and customizing them with users' own data.\n",
    "\n",
    "<img align=\"center\" src=\"https://d29g4g2dyqv443.cloudfront.net/sites/default/files/akamai/TAO/tlt-tao-toolkit-bring-your-own-model-diagram.png\" width=\"1080\"> "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Learning Objectives\n",
    "\n",
    "In this notebook, you will learn how to leverage the simplicity and convenience of TAO to:\n",
    "\n",
    "* Take a pretrained model and train a GestureNet model on HGR dataset\n",
    "* Run Inference on the trained model\n",
    "* Export the retrained model to a .etlt file for deployment to DeepStream SDK\n",
    "\n",
    "At the end of this notebook, you will have generated a trained and optimized `gesturenet` model, \n",
    "which you may deploy via [DeepStream](https://developer.nvidia.com/deepstream-sdk).\n",
    "\n",
    "### Table of Contents\n",
    "\n",
    "This notebook shows an example of classifying gestures using GestureNet in the Train Adapt Optimize (TAO) Toolkit.\n",
    "\n",
    "0. [Set up env variables, map drives, and install dependencies](#head-0)\n",
    "1. [Install the TAO launcher](#head-1)\n",
    "2. [Prepare dataset and pre-trained model](#head-2) <br>\n",
    "    A. [Verify and prepare dataset](#head-2-1) <br>\n",
    "    B. [Generate hand crops and dataset json](#head-2-2) <br>\n",
    "    C. [Download pre-trained model](#head-2-3) <br>\n",
    "3. [Provide training specification](#head-3) <br>\n",
    "4. [Run TAO training](#head-4) <br>\n",
    "5. [Evaluate the trained model](#head-5) <br>\n",
    "6. [Export](#head-6) <br>\n",
    "7. [Inference](#head-7) <br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. Set up env variables, map drives and install dependencies <a class=\"anchor\" id=\"head-0\"></a>\n",
    "When using the purpose-built pretrained models from NGC, please make sure to set the `$KEY` environment variable to the key as mentioned in the model overview. Failing to do so, can lead to errors when trying to load them as pretrained models.\n",
    "\n",
    "The following notebook requires the user to set an env variable called the `$LOCAL_PROJECT_DIR` as the path to the users' workspace. Please note that the dataset to run this notebook is expected to reside in the `$LOCAL_PROJECT_DIR/gesturenet/data`, while the TAO experiment generated collaterals will be output to `$LOCAL_PROJECT_DIR/gesturenet`. More information on how to set up the dataset and the supported steps in the TAO workflow are provided in the subsequent cells.\n",
    "\n",
    "*Note: This notebook currently is by default set up to run training using 1 GPU. To use more GPU's please update the env variable `$NUM_GPUS` accordingly*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setting up env variables for cleaner command-line commands.\n",
    "import os\n",
    "\n",
    "%env KEY=nvidia_tlt\n",
    "%env NUM_GPUS=1\n",
    "%env USER_EXPERIMENT_DIR=/workspace/tao-experiments/gesturenet\n",
    "%env DATA_DOWNLOAD_DIR=/workspace/tao-experiments/gesturenet/data\n",
    "\n",
    "# Set this path if you don't run the notebook from the samples directory.\n",
    "# %env NOTEBOOK_ROOT=~/tao-samples/gesturenet\n",
    "\n",
    "# Please define this local project directory that needs to be mapped to the TAO docker session.\n",
    "# The dataset is expected to be present in $LOCAL_PROJECT_DIR/gesturenet/data, while the results for the steps\n",
    "# in this notebook will be stored at $LOCAL_PROJECT_DIR/gesturenet\n",
    "# !PLEASE MAKE SURE TO UPDATE THIS PATH!.\n",
    "%env LOCAL_PROJECT_DIR=/path/to/local/experiments\n",
    "\n",
    "# $PROJECT_DIR is the path to the sample notebook folder and the dependency folder\n",
    "# $PROJECT_DIR/deps should exist for dependency installation\n",
    "%env PROJECT_DIR=/path/to/local/samples_dir\n",
    "\n",
    "os.environ[\"LOCAL_DATA_DIR\"] = os.path.join(\n",
    "    os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()),\n",
    "    \"gesturenet/data\"\n",
    ")\n",
    "os.environ[\"LOCAL_EXPERIMENT_DIR\"] = os.path.join(\n",
    "    os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()),\n",
    "    \"gesturenet\"\n",
    ")\n",
    "\n",
    "# The sample spec files are present in the same path as the downloaded samples.\n",
    "os.environ[\"LOCAL_SPECS_DIR\"] = os.path.join(\n",
    "    os.getenv(\"NOTEBOOK_ROOT\", os.getcwd()),\n",
    "    \"specs\"\n",
    ")\n",
    "\n",
    "%env SPECS_DIR=/workspace/tao-experiments/gesturenet/specs\n",
    "\n",
    "# Showing list of specification files.\n",
    "!ls -rlt $LOCAL_SPECS_DIR"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The cell below maps the project directory on your local host to a workspace directory in the TAO docker instance, so that the data and the results are mapped from in and out of the docker. For more information please refer to the [launcher instance](https://docs.nvidia.com/tao/tao-toolkit/text/tao_launcher.html) in the user guide.\n",
    "\n",
    "When running this cell on AWS, update the drive_map entry with the dictionary defined below, so that you don't have permission issues when writing data into folders created by the TAO docker.\n",
    "\n",
    "```json\n",
    "drive_map = {\n",
    "    \"Mounts\": [\n",
    "            # Mapping the data directory\n",
    "            {\n",
    "                \"source\": os.environ[\"LOCAL_PROJECT_DIR\"],\n",
    "                \"destination\": \"/workspace/tao-experiments\"\n",
    "            },\n",
    "            # Mapping the specs directory.\n",
    "            {\n",
    "                \"source\": os.environ[\"LOCAL_SPECS_DIR\"],\n",
    "                \"destination\": os.environ[\"SPECS_DIR\"]\n",
    "            },\n",
    "        ],\n",
    "    \"DockerOptions\": {\n",
    "        \"user\": \"{}:{}\".format(os.getuid(), os.getgid())\n",
    "    }\n",
    "}\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Mapping up the local directories to the TAO docker.\n",
    "import json\n",
    "mounts_file = os.path.expanduser(\"~/.tao_mounts.json\")\n",
    "\n",
    "# Define the dictionary with the mapped drives\n",
    "drive_map = {\n",
    "    \"Mounts\": [\n",
    "        # Mapping the data directory\n",
    "        {\n",
    "            \"source\": os.environ[\"LOCAL_PROJECT_DIR\"],\n",
    "            \"destination\": \"/workspace/tao-experiments\"\n",
    "        },\n",
    "        # Mapping the specs directory.\n",
    "        {\n",
    "            \"source\": os.environ[\"LOCAL_SPECS_DIR\"],\n",
    "            \"destination\": os.environ[\"SPECS_DIR\"]\n",
    "        },\n",
    "    ]\n",
    "}\n",
    "\n",
    "# Writing the mounts file.\n",
    "with open(mounts_file, \"w\") as mfile:\n",
    "    json.dump(drive_map, mfile, indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat ~/.tao_mounts.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install requirement\n",
    "!pip3 install Cython==0.29.36\n",
    "!pip3 install -r $PROJECT_DIR/deps/requirements-pip.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Install the TAO launcher <a class=\"anchor\" id=\"head-1\"></a>\n",
    "The TAO launcher is a python package distributed as a python wheel listed in PyPI. You may install the launcher by executing the following cell.\n",
    "\n",
    "Please note that TAO Toolkit recommends users run the TAO launcher in a virtual env with python 3.6.9. You may follow the instruction on this [page](https://virtualenvwrapper.readthedocs.io/en/latest/install.html) to set up a python virtual env using the `virtualenv` and `virtualenvwrapper` packages. Once you have set up virtualenvwrapper, please set the version of python to be used in the virtual env by using the `VIRTUALENVWRAPPER_PYTHON` variable. You may do so by running\n",
    "\n",
    "```sh\n",
    "export VIRTUALENVWRAPPER_PYTHON=/path/to/bin/python3.x\n",
    "```\n",
    "where x >= 6 and <= 8\n",
    "\n",
    "We recommend performing this step first and then launching the notebook from the virtual environment. In addition to installing TAO python package, please make sure of the following software requirements:\n",
    "* python >=3.7, <=3.10.x\n",
    "* docker-ce > 19.03.5\n",
    "* docker-API 1.40\n",
    "* nvidia-container-toolkit > 1.3.0-1\n",
    "* nvidia-container-runtime > 3.4.0-1\n",
    "* nvidia-docker2 > 2.5.0-1\n",
    "* nvidia-driver > 455+\n",
    "\n",
    "Once you have installed the pre-requisites, please log in to the docker registry nvcr.io by following the command below\n",
    "\n",
    "```sh\n",
    "docker login nvcr.io\n",
    "```\n",
    "\n",
    "You will be triggered to enter a username and password. The username is `$oauthtoken` and the password is the API key generated from `ngc.nvidia.com`. Please follow the instructions in the [NGC setup guide](https://docs.nvidia.com/ngc/ngc-overview/index.html#generating-api-key) to generate your own API key."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Skip this cell if the TAO launcher was already installed.\n",
    "!pip3 install nvidia-tao"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# View the version of the TAO launcher\n",
    "!tao info"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Prepare dataset and pre-trained model <a class=\"anchor\" id=\"head-2\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will be using the database for hand gesture recognition (HGR) for the tutorial. To find more details, please visit http://sun.aei.polsl.pl/~mkawulok/gestures/. Please download the HGR1 [images](http://sun.aei.polsl.pl/~mkawulok/gestures/hgr1_images.zip), [feature points](http://sun.aei.polsl.pl/~mkawulok/gestures/hgr1_feature_pts.zip) and HGR2B [images](http://sun.aei.polsl.pl/~mkawulok/gestures/hgr2b_images.zip), [feature points](http://sun.aei.polsl.pl/~mkawulok/gestures/hgr2b_feature_pts.zip) and place the zip files in `$LOCAL_DATA_DIR`. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A. Verify and prepare dataset <a class=\"anchor\" id=\"head-2-1\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check the zip files are present.\n",
    "!mkdir -p $LOCAL_DATA_DIR\n",
    "!if [ ! -f $LOCAL_DATA_DIR/hgr1_images.zip ]; then echo 'hgr1_images zip file not found, please download.'; else echo 'Found hgr1_images zip file.';fi\n",
    "!if [ ! -f $LOCAL_DATA_DIR/hgr1_feature_pts.zip ]; then echo 'hgr1_feature_pts zip file not found, please download.'; else echo 'Found hgr1_feature_pts zip file.';fi\n",
    "!if [ ! -f $LOCAL_DATA_DIR/hgr2b_images.zip ]; then echo 'hgr2b_images zip file not found, please download.'; else echo 'Found hgr2b_images zip file.';fi\n",
    "!if [ ! -f $LOCAL_DATA_DIR/hgr2b_feature_pts.zip ]; then echo 'hgr2b_feature_pts zip file not found, please download.'; else echo 'Found hgr2b_feature_pts zip file.';fi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# unpack downloaded datasets to $DATA_DOWNLOAD_DIR.\n",
    "# The images will be under $DATA_DOWNLOAD_DIR/original_images and $DATA_DOWNLOAD_DIR/feature_points\n",
    "!unzip -u ${LOCAL_DATA_DIR}/hgr1_images.zip -d ${LOCAL_DATA_DIR}\n",
    "!unzip -u ${LOCAL_DATA_DIR}/hgr1_feature_pts.zip -d ${LOCAL_DATA_DIR}\n",
    "!unzip -u ${LOCAL_DATA_DIR}/hgr2b_images.zip -d ${LOCAL_DATA_DIR}\n",
    "!unzip -u ${LOCAL_DATA_DIR}/hgr2b_feature_pts.zip -d ${LOCAL_DATA_DIR}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert dataset to required format for gesturenet dataset_convert\n",
    "!python3 convert_hgr_to_tlt_data.py --input_image_dir=$LOCAL_DATA_DIR/original_images \\\n",
    "                                    --input_label_file=$LOCAL_DATA_DIR/feature_points \\\n",
    "                                    --output_dir=$LOCAL_EXPERIMENT_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# verify\n",
    "import os\n",
    "\n",
    "LOCAL_EXPERIMENT_DIR = os.environ.get('LOCAL_EXPERIMENT_DIR')\n",
    "num_labels = len(os.listdir(os.path.join(LOCAL_EXPERIMENT_DIR, \"original/data/annotation\")))\n",
    "print(\"Number of labels in the dataset. {}\".format(num_labels))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### B. Generate hand crops and dataset json <a class=\"anchor\" id=\"head-2-2\"></a>\n",
    "\n",
    "* Update the `dataset_config.json` and `dataset_experiment_config.json` spec files\n",
    "* Create the crop and json using the gesturenet dataset_convert \n",
    "\n",
    "*Note: Crops and dataset json only need to be generated once.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Hand crop generation spec file\")\n",
    "!cat $LOCAL_SPECS_DIR/dataset_config.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Dataset experiment spec file\")\n",
    "!cat $LOCAL_SPECS_DIR/dataset_experiment_config.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model gesturenet dataset_convert --dataset_spec $SPECS_DIR/dataset_config.json \\\n",
    "                                --k_folds 0 \\\n",
    "                                --experiment_spec $SPECS_DIR/dataset_experiment_config.json \\\n",
    "                                --output_filename $USER_EXPERIMENT_DIR/data.json \\\n",
    "                                --experiment_name v1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check to see if proper json file is generated.\n",
    "!if [ ! -f $LOCAL_EXPERIMENT_DIR/data.json ]; then echo \"Json file was not generated properly.\"; else echo \"Json was generated properly.\"; fi"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### C. Download pre-trained model <a class=\"anchor\" id=\"head-2-3\"></a>\n",
    "\n",
    "Please follow the instructions in the following to download and verify the pretrained model for gesturenet.\n",
    "\n",
    "For FpeNet pretrained model please download model: `nvidia/tao/gesturenet:trainable_v1.0`.\n",
    "\n",
    "After obtaining the pre-trained model, please place the model in $LOCAL_EXPERIMENT_DIR\n",
    "\n",
    "You will have the following path-\n",
    "\n",
    "* pretrained model in `$LOCAL_EXPERIMENT_DIR/pretrained_models/gesturenet_vtrainable_v1.0/model.tlt`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Installing NGC CLI on the local machine.\n",
    "## Download and install\n",
    "%env CLI=ngccli_cat_linux.zip\n",
    "!mkdir -p $LOCAL_PROJECT_DIR/ngccli\n",
    "\n",
    "# Remove any previously existing CLI installations\n",
    "!rm -rf $LOCAL_PROJECT_DIR/ngccli/*\n",
    "!wget \"https://ngc.nvidia.com/downloads/$CLI\" -P $LOCAL_PROJECT_DIR/ngccli\n",
    "!unzip -u \"$LOCAL_PROJECT_DIR/ngccli/$CLI\" -d $LOCAL_PROJECT_DIR/ngccli/\n",
    "!rm $LOCAL_PROJECT_DIR/ngccli/*.zip \n",
    "os.environ[\"PATH\"]=\"{}/ngccli/ngc-cli:{}\".format(os.getenv(\"LOCAL_PROJECT_DIR\", \"\"), os.getenv(\"PATH\", \"\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# List models available in the model registry.\n",
    "!ngc registry model list nvidia/tao/gesturenet:*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the target destination to download the model.\n",
    "!mkdir -p $LOCAL_EXPERIMENT_DIR/pretrained_models/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download the pretrained model from NGC\n",
    "!ngc registry model download-version nvidia/tao/gesturenet:trainable_v1.0 \\\n",
    "    --dest $LOCAL_EXPERIMENT_DIR/pretrained_models/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls -rlt $LOCAL_EXPERIMENT_DIR/pretrained_models/gesturenet_vtrainable_v1.0 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check the model is present\n",
    "!if [ ! -f $LOCAL_EXPERIMENT_DIR/pretrained_models/gesturenet_vtrainable_v1.0/model.tlt ]; then echo 'Pretrained model file not found, please download.'; else echo 'Found Pretrain model file.';fi"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Provide training specification <a class=\"anchor\" id=\"head-3\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* Dataset configuration\n",
    "    * In order to load the data properly, you will need to change the `dataset:data_path` to the generated `json` (folder and file) file generated in part B above. By default it is located at `$LOCAL_SPECS_DIR/data.json`\n",
    "    * Update number of classes and class number to name map\n",
    "* Pre-trained models. There is an optional parameter to load head of model. Only set `add_new_head: false` if you want to finetune on dataset with same gestures as pretrained model. Please ensure the gesture class to index map matches pretrained model.\n",
    "* Augmentation parameters for on the fly data augmentation\n",
    "* Other training (hyper-)parameters such as batch size, number of epochs, learning rate etc.|"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat $LOCAL_SPECS_DIR/train_spec.json"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Run TAO training <a class=\"anchor\" id=\"head-4\"></a>\n",
    "* Provide the sample spec file and the encryption key"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model gesturenet train -e $SPECS_DIR/train_spec.json \\\n",
    "                      -k $KEY"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Evaluate the trained model <a class=\"anchor\" id=\"head-5\"></a>\n",
    "\n",
    "* Please update model path to location trained model is saved at"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model gesturenet evaluate -e $USER_EXPERIMENT_DIR/model/train_spec.json \\\n",
    "                         -m $USER_EXPERIMENT_DIR/model/model.tlt \\\n",
    "                         -k $KEY"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Inference <a class=\"anchor\" id=\"head-7\"></a>\n",
    "In this section, we run the `gesturenet inference` tool to generate inferences on the trained models. Please ensure the spec file `inference.json` is configured correctly. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model gesturenet inference -e $USER_EXPERIMENT_DIR/model/train_spec.json \\\n",
    "                          -m $USER_EXPERIMENT_DIR/model/model.tlt \\\n",
    "                          -k $KEY \\\n",
    "                          --image_root_path $USER_EXPERIMENT_DIR \\\n",
    "                          --data_json $USER_EXPERIMENT_DIR/data.json \\\n",
    "                          --data_type kpi_set \\\n",
    "                          --results_dir $USER_EXPERIMENT_DIR/model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import cv2\n",
    "import IPython.display\n",
    "import json\n",
    "import PIL.Image\n",
    "\n",
    "json_spec_path = os.path.join(os.environ.get('LOCAL_EXPERIMENT_DIR'), 'data.json')\n",
    "data_type = \"kpi_set\"\n",
    "result_file = os.path.join(os.environ.get('LOCAL_EXPERIMENT_DIR'), 'model/results.txt')\n",
    "model_spec_path = os.path.join(os.environ.get('LOCAL_EXPERIMENT_DIR'), 'model/train_spec.json')\n",
    "\n",
    "# Read in json spec.\n",
    "with open(json_spec_path, 'r') as file:\n",
    "    full_spec = json.load(file)\n",
    "spec = full_spec[data_type]\n",
    "\n",
    "# Read in model spec.\n",
    "with open(model_spec_path, 'r') as file:\n",
    "    model_spec = json.load(file)\n",
    "\n",
    "class_labels = model_spec['dataset']['classes']\n",
    "\n",
    "results = open(result_file, 'r')\n",
    "\n",
    "images = spec['images']\n",
    "\n",
    "for image_dict in images:\n",
    "\n",
    "    image_path = os.path.join(os.environ.get('LOCAL_EXPERIMENT_DIR'), image_dict['full_image_path'])\n",
    "    bbox = image_dict['bbox_coordinates']\n",
    "    image = cv2.imread(image_path)\n",
    "    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
    "    # Get corners of rectangle.\n",
    "    upper_left = tuple(bbox[0])\n",
    "    bottom_right = tuple(bbox[3])\n",
    "    # draw rectangle onto image.\n",
    "    cv2.rectangle(image, upper_left, bottom_right, (0, 255, 0), 2)\n",
    "\n",
    "    if image is None:\n",
    "        results.readline()\n",
    "        continue\n",
    "    image_result = results.readline()\n",
    "    prediction = image_result.split(' ')[1]\n",
    "    # Get class label.\n",
    "    label = list(class_labels.keys())[list(class_labels.values()).index(int(prediction))]\n",
    "    # Get bottom right corner.\n",
    "    x = 0\n",
    "    y = image.shape[0]-5\n",
    "    # Display Image.\n",
    "    image = cv2.putText(image, label, (x, y), cv2.FONT_HERSHEY_COMPLEX_SMALL, 1, (255, 24, 8))\n",
    "    IPython.display.display(PIL.Image.fromarray(image))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Export <a class=\"anchor\" id=\"head-6\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 7.1 Export .etlt model\n",
    "\n",
    "Use the export functionality to export an encrypted model in fp32 format without any optimizations.\n",
    "\n",
    "* Modify `-m` to your model path\n",
    "* Modify `--out_file` to your desired output path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model gesturenet export -m $USER_EXPERIMENT_DIR/model/model.tlt \\\n",
    "                       -k ${KEY} \\\n",
    "                       --out_file $USER_EXPERIMENT_DIR/model/model.etlt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# check the deployment file is presented\n",
    "!if [ ! -f $LOCAL_EXPERIMENT_DIR/model/model.etlt ]; then echo 'Deployment file not found, please generate.'; else echo 'Found deployment file.';fi"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 7.2 INT8 Optimization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "GestureNet model supports int8 inference mode in TensorRT. In order to do this, the model is first calibrated to run 8-bit inferences. This is the process:\n",
    "\n",
    "* Provide a directory with set of images to be used for calibration.\n",
    "* A calibration tensorfile is generated and saved in --cal_data_file\n",
    "* This tensorfile is use to calibrate the model and the calibration table is stored in --cal_cache_file\n",
    "* The calibration table in addition to the model is used to generate the int8 tensorrt engine to the path --engine_file\n",
    "\n",
    "*Note: For this example, we generate a calibration tensorfile containing 100 batches of training data. Ideally, it is best to use at least 10-20% of the training data to do so. The more data provided during calibration, the closer int8 inferences are to fp32 inferences.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Number of calibration samples to use\n",
    "%set_env NUM_CALIB_SAMPLES=100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python3 sample_calibration_images.py \\\n",
    "    -a $LOCAL_EXPERIMENT_DIR/data.json \\\n",
    "    -i $LOCAL_EXPERIMENT_DIR \\\n",
    "    -o $LOCAL_EXPERIMENT_DIR/calibration_samples/ \\\n",
    "    -n $NUM_CALIB_SAMPLES \\\n",
    "    --randomize"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 7.3 Export Deployable INT8 Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model gesturenet export -m $USER_EXPERIMENT_DIR/model/model.tlt \\\n",
    "                       -k $KEY \\\n",
    "                       --engine_file $USER_EXPERIMENT_DIR/model/model.int8.engine \\\n",
    "                       --data_type int8 \\\n",
    "                       --cal_image_dir $USER_EXPERIMENT_DIR/calibration_samples/ \\\n",
    "                       --cal_cache_file $USER_EXPERIMENT_DIR/model/int8_calibration.bin \\\n",
    "                       --cal_data_file $USER_EXPERIMENT_DIR/model/int8_calibration.tensorfile \\\n",
    "                       --batches 100"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 7.4 Run Inference on Exported INT8 Engine File"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model gesturenet inference -e $USER_EXPERIMENT_DIR/model/train_spec.json \\\n",
    "                          -m $USER_EXPERIMENT_DIR/model/model.int8.engine \\\n",
    "                          -k $KEY \\\n",
    "                          --image_root_path $USER_EXPERIMENT_DIR \\\n",
    "                          --data_json $USER_EXPERIMENT_DIR/data.json \\\n",
    "                          --data_type kpi_set \\\n",
    "                          --results_dir $USER_EXPERIMENT_DIR/model/model_int8_engine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# check the results file is generated\n",
    "!if [ ! -f $LOCAL_EXPERIMENT_DIR/model/model_int8_engine/results.txt ]; then echo 'Results file not found!'; else cat $LOCAL_EXPERIMENT_DIR/model/model_int8_engine/results.txt;fi"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.10 64-bit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  },
  "vscode": {
   "interpreter": {
    "hash": "767d51c1340bd893661ea55ea3124f6de3c7a262a8b4abca0554b478b1e2ff90"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
